# $Id: flowlib.py 83306ac00e39 2013/03/03 02:04:53 pgurumur $
# Copyright (c) 2009 Prabhu Gurumurthy  <pgurumur@gmail.com>
# 
# Permission to use, copy, modify, and distribute this software for any
# purpose with or without fee is hereby granted, provided that the above
# copyright notice and this permission notice appear in all copies.
#
# THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
# WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
# MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
# ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
# WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
# ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
# OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
#
__version__  = "$Revision: 83306ac00e39 $"
__author__   = "$Author: pgurumur $"
__modified__ = "$Date: 2013-03-03 02:04:53Z $"

import warnings
import struct
import collections
import datetime

from venya.lib import stdlib, ip

def parseNetFlowData(packet, start, template, length):
   end = start + 4
   flowitems = []
   while end < length:
      flow = collections.OrderedDict({})
      for (key, value) in template.items():
         pktformat = None
         if value == 4:
            pktformat = "!I"
         elif value == 2:
            pktformat = "!H"
            end -= 2
         elif value == 1:
            pktformat = "!c"
            end -= 3
         elif value == 6:
            pktformat = "!HHH"
            end += 2

         data = None
         if pktformat:
            if value == 6:
               data = struct.unpack(pktformat, packet[start:end])
            else:
               data = struct.unpack(pktformat, packet[start:end])[0]

         if data:
            if isinstance(data, str):
               data = ord(data)

            if key in ('IPV4_SRC_ADDR',
                  'IPV4_DST_ADDR', 'IPV4_NEXT_HOP') and data:
               data = ip.ip(data, version = 4)

            if key in ('IN_SRC_MAC', 'OUT_DST_MAC') and data:
               data = "%x.%x.%x" %(data[0], data[1], data[2])

            flow[key] = data

         start = end
         end += 4

      if flow:
         flowitems.append(flow)

   return flowitems

template = {
      1: 'IN_BYTES',
      2: 'IN_PKTS',
      3: 'FLOWS',
      4: 'PROTOCOL',
      5: 'SRC_TOS',
      6: 'TCP_FLAGS',
      7: 'L4_SRC_PORT',
      8: 'IPV4_SRC_ADDR',
      9: 'SRC_MASK',
      10: 'INPUT_SNMP',
      11: 'L4_DST_PORT',
      12: 'IPV4_DST_ADDR',
      13: 'DST_MASK',
      14: 'OUTPUT_SNMP',
      15: 'IPV4_NEXT_HOP',
      16: 'SRC_AS',
      17: 'DST_AS',
      18: 'BGP_IPV4_NEXT_HOP',
      19: 'MUL_DST_PKTS',
      20: 'MUL_DST_BYTES',
      21: 'LAST_SWITCHED',
      22: 'FIRST_SWITCHED',
      23: 'OUT_BYTES',
      24: 'OUT_PKTS',
      25: 'MIN_PKT_LENGTH',
      26: 'MAX_PKT_LENGTH',
      27: 'IPV6_SRC_ADDR',
      28: 'IPV6_DST_ADDR',
      29: 'IPV6_SRC_MASK',
      27: 'IPV6_DST_MASK',
      31: 'IPV6_FLOW_LABEL',
      32: 'ICMP_TYPE',
      33: 'MUL_IGMP_TYPE',
      34: 'SAMPLING_INTERVAL',
      35: 'SAMPLING_ALGORITHM',
      36: 'FLOW_ACTIVE_TIMEOUT',
      37: 'FLOW_INACTIVE_TIMEOUT',
      38: 'ENGINE_TYPE',
      39: 'ENGINE_ID',
      40: 'TOTAL_BYTES_EXP',
      41: 'TOTAL_PKTS_EXP',
      42: 'TOTAL_FLOWS_EXP',
      43: 'PROPRIETARY',
      44: 'IPV4_SRC_PREFIX',
      45: 'IPV4_DST_PREFIX',
      46: 'MPLS_TOP_LABEL_TYPE',
      47: 'MPLS_TOP_LABEL_IPADDR',
      48: 'FLOW_SAMPLER_ID',
      49: 'FLOW_SAMPLER_MODE',
      50: 'FLOW_SAMPLER_RANDOM_INTERVAL',
      51: 'PROPRIETARY',
      52: 'MIN_TTL',
      53: 'MAX_TTL',
      54: 'IPV4_IDENT',
      55: 'DST_TOS',
      56: 'IN_SRC_MAC',
      57: 'OUT_DST_MAC',
      58: 'SRC_VLAN',
      59: 'DST_VLAN',
      60: 'IP_VERSION',
      61: 'DIRECTION',
      62: 'IPV6_NEXT_HOP',
      63: 'BGP_IPV6_NEXT_HOP',
      64: 'IPV6_OPTION_HEADERS',
      70: 'MPLS_LABEL1',
      71: 'MPLS_LABEL2',
      72: 'MPLS_LABEL3',
      73: 'MPLS_LABEL4',
      74: 'MPLS_LABEL5',
      75: 'MPLS_LABEL6',
      76: 'MPLS_LABEL7',
      77: 'MPLS_LABEL8',
      78: 'MPLS_LABEL9',
      79: 'MPLS_LABEL10',
      80: 'IN_DST_MAC',
      81: 'OUT_SRC_MAC',
      82: 'IF_NAME',
      83: 'IF_DESC',
      84: 'SAMPLER_NAME',
      85: 'IN_PERM_BYTES',
      86: 'IN_PERM_PKTS',
      88: 'FRAGMENT_OFFSET',
      89: 'FORWARDING_STATUS',
      90: 'MPLS_PAL_RD',
      91: 'MPLS_PREFIX_LEN',
      92: 'SRC_TRAFFIC_INDEX',
      93: 'DST_TRAFFIC_INDEX',
      94: 'APP_DESCR',
      95: 'APP_TAG',
      96: 'APP_NAME',
      98: 'DIFFSERV_PT',
      99: 'REPLICATION_FACTOR',
      102: 'L2_PKT_SECTION_OFFSET',
      103: 'L2_PKT_SECTION_SIZE',
      104: 'L2_PKT_SECTION_DATA'
      }

class NetFlowHeader(object):
   def __init__(self, Packet):
      super(NetFlowHeader, self).__init__()
      self._packet = Packet
      self._version = 0
      self._count = 0
      self._uptime = 0
      self._unixtime = 0
      self._sequence = 0
      self._srcid = 0
      self._flowtype = -1
      self._templateFlowSet = None
      self._dataFlowSet = None
      self._optionsTemplateFlowSet = None

      header = None
      try:
         header = struct.unpack("!HHIIII", self._packet[:20])
      except struct.error as err:
         raise ValueError(err)

      try:
         self._version = int(header[0])
         self._count = int(header[1])
         self._uptime = header[2] / 1000
         self._unixtime = int(header[3])
         self._sequence = int(header[4])
         self._srcid = int(header[5])
      except ValueError as err:
         self.error(str(err))
      except TypeError as err:
         self.error(str(err))

      try:
         if len(self._packet) >= 22:
            self._flowtype = struct.unpack("!H", self._packet[20:22])[0]
      except TypeError as err:
         self.error(str(err))
      except struct.error as err:
         self.error(str(err))

   def __int__(self):
      return self.sequence

   def __len__(self):
      return self.length

   def AddTemplateFlowSet(self, Packet):
      self._templateFlowSet = TemplateFlowSet(Packet)

   def AddDataFlowSet(self, Packet):
      self._dataFlowSet = DataFlowSet(Packet)

   def AddOptionsTemplateSet(self, Packet):
      self._optionsTemplateFlowSet = OptionsTemplateFlowSet(Packet)

   def error(self, fmt, *args):
      warnings.warn(stdlib.sprintf(fmt, *args))

   version = property(fget = lambda self: self._version)
   sequence = property(fget = lambda self: self._sequence)
   count = property(fget = lambda self: self._count)
   length = property(fget = lambda self: self.count)
   flowtype = property(fget = lambda self: self._flowtype)

class TemplateFlowSet(object):
   def __init__(self, Packet):
      self._packet = None
      self._fid = -1
      self._length = 0
      self._tid = 0
      self._count = 0
      if Packet and len(Packet):
         self._packet = Packet

      self._template = collections.OrderedDict({})
      self._flowdata = []

      try:
         (self._fid, self._length, self._tid, self._count) = struct.unpack(
               "!HHHH", self._packet[20:28])
      except struct.error as err:
         self.error(str(err))
      except TypeError as err:
         self.error(str(err))

      print(self._fid, self._length, self._tid, self._count)

      start = 28
      end = 32
      for ix in xrange(0, int(self._count)):
         (value, size) = struct.unpack("!HH", self._packet[start:end])
         start = end
         end += 4
         self._template[template[int(value)]] = int(size)

      sequence, length = struct.unpack("!HH", self._packet[start:end])
      self._flowdata = parseNetFlowData(
            self._packet, end, self._template, length)

   def __int__(self):
      return self._tid

   def __len__(self):
      return self.length

   length = property(fget = lambda self: self._length)
   fsetid = property(fget = lambda self: self._fid)
   template = property(fget = lambda self:  self._template)
   data = property(fget = lambda self: self._flowdata)

class DataFlowSet(object):
   def __init__(self, Packet, TemplateId = 0, Template = None):
      self._packet = None
      self._fid = -1
      self._length = 0

      if Packet and len(Packet):
         self._packet = Packet

      try:
         (self._fid, self._length) = struct.unpack("!HH", self._packet[20:24])
      except struct.error as err:
         raise ValueError(err)

      print(self._fid, self._length, TemplateId, Template)

   def __len__(self):
      return self.length

   length = property(fget = lambda self: self._length)
   fsetid = property(fget = lambda self: self._fid)
